import os
import tensorflow as tf
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.compat.v1 import ConfigProto
from tensorflow.compat.v1 import InteractiveSession
# from Tensorflow2.hw3.Model import *
import matplotlib.pyplot as plt
from Core.ResNet_Model import *
from Core.dataset import *
import pandas as pd

'''
改用ResNet model
'''
if __name__ == '__main__':
    # 设置GPU显存自适应
    config = ConfigProto()
    config.gpu_options.allow_growth = True
    session = InteractiveSession(config=config)

    # get dataset
    train_generator, validation_generator, _ = get_generator()

    # model params config
    # hyper_params
    lr = 1e-4
    optimizer = tf.keras.optimizers.Adam(lr)
    loss = tf.keras.losses.CategoricalCrossentropy()
    epochs = 100
    # batch_size 在generator里面定义
    # batch_size = 32

    # ResNet18 Model
    model = ResNet18([2, 2, 2, 2])
    # checkpoint
    saved_path = r'checkpoint_ResNet'
    if not os.path.exists(saved_path):
        os.mkdir(saved_path)
    checkpoint_saved_path = os.path.join(saved_path, 'CovNet.ckpt')
    if os.path.exists(checkpoint_saved_path + '.index'):
        print('--load exist model--')
        model.load_weights(checkpoint_saved_path)
    callbacks = [
        tf.keras.callbacks.ModelCheckpoint(filepath=checkpoint_saved_path,
                                           monitor='val_loss',
                                           save_best_only=True,
                                           save_weights_only=True),
        tf.keras.callbacks.TensorBoard(log_dir=saved_path),
        tf.keras.callbacks.EarlyStopping(monitor='loss',
                                         patience=10,
                                         min_delta=0.01,
                                         restore_best_weights=True)
    ]
    model.compile(optimizer=optimizer, loss=loss, metrics=['acc'])
    history = model.fit(train_generator,
                        # batch_size=batch_size,
                        epochs=epochs,
                        validation_data=validation_generator,
                        callbacks=callbacks,
                        shuffle=True)
    model.summary()

    pd.DataFrame(history.history).plot()
    plt.grid = True
    plt.gca().set_ylim(0, 1)
    plt.show()


